# -*- coding: utf-8 -*-
"""
Created on Sun Feb 28 21:37:22 2021

@author: xiaoyifu
"""

import numpy as np
import sys
sys.path.insert(0, "../")
from ge.classify import read_node_label, Classifier
from ge import Node2Vec
from sklearn.linear_model import LogisticRegression

import matplotlib.pyplot as plt
import networkx as nx
from sklearn.manifold import TSNE


def evaluate_embeddings(embeddings):
    X, Y = read_node_label('../../data/1/rna/rna(weight).edgelist')#./data/miRNA/miRNA_label.txt
    tr_frac = 0.8
    print("Training classifier using {:.2f}% nodes...".format(
        tr_frac * 100))
    clf = Classifier(embeddings=embeddings, clf=LogisticRegression())
    clf.split_train_evaluate(X, Y, tr_frac)


def plot_embeddings(embeddings,):
    X, Y = read_node_label('../data/RNA/RNA_label.txt')#./data/mRNA-lncRNA/mRNA_lncRNA_label.txt

    emb_list = []
    for k in X:
        emb_list.append(embeddings[k])
    emb_list = np.array(emb_list)

    model = TSNE(n_components=2)
    node_pos = model.fit_transform(emb_list)

    color_idx = {}
    for i in range(len(X)):
        color_idx.setdefault(Y[i][0], [])
        color_idx[Y[i][0]].append(i)

    for c, idx in color_idx.items():
        plt.scatter(node_pos[idx, 0], node_pos[idx, 1], label=c)
    plt.legend()
    plt.show()


if __name__ == "__main__":
    G = nx.read_edgelist('../../../data/rna/1/rna(weight).edgelist',
                         create_using = nx.DiGraph(), nodetype =None, data = [('weight',float)])#read graph

    model = Node2Vec(G, walk_length=10, num_walks=80,
                     p=10, q=0.5, workers=2, use_rejection_sampling=0)
    model.train(window_size = 5, iter = 2,embed_size=512)
    embeddings = model.get_embeddings()# get embedding vectors

    #evaluate_embeddings(embeddings)
    #plot_embeddings(embeddings)
    np.save("./node2vec/miRNA_embedding.npy",embeddings)#mRNA_lncRNA_embedding.npy miRNA_embedding1.npy

